package edu.washington.ling.syntric.proc;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.configuration.ConfigurationException;
import org.apache.commons.configuration.XMLConfiguration;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.washington.ling.syntric.config.Configuration;
import edu.washington.ling.syntric.io.ExampleDoc;
import edu.washington.ling.syntric.io.ExampleDocReader;
import edu.washington.ling.syntric.io.Example;
import edu.washington.ling.syntric.io.RuleAndExampleFileWriter;
import edu.washington.ling.syntric.io.RuleFileWriter;
import edu.washington.ling.syntric.io.RuleWriter;
import edu.washington.ling.syntric.io.TreeReader;
import edu.washington.ling.syntric.struct.error.*;
import edu.washington.ling.syntric.struct.rule.UniDSPattern;
import edu.washington.ling.syntric.struct.rule.UniPSPattern;
import edu.washington.ling.syntric.struct.rule.UniRule;
import edu.washington.ling.syntric.struct.tree.Token;
import edu.washington.ling.syntric.struct.tree.Tree;
import edu.washington.ling.syntric.util.ConvertUtils;
import edu.washington.ling.syntric.util.MergeUtils;

public class ExtractUnaryRules {
	
	private static ExtractUnaryRulesOptions opts = new ExtractUnaryRulesOptions();
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
		
	private static void printRuleStatistics(Map<UniDSPattern, List<UniPSPattern>> ruleMap) {
		int numAmbig = 0;
		int numBiAmbig = 0;
		int numTriAmbig = 0;
		int numMoreAmbig = 0;
		for (UniDSPattern dsPatt: ruleMap.keySet()) {
			Integer count = ruleMap.get(dsPatt).size();
			if (count > 1) {
				numAmbig += 1;
			}
			if (count == 2) {
				numBiAmbig += 1;
			}
			else if (count == 3) {
				numTriAmbig += 1;
			}
			else if (count > 3) {
				numMoreAmbig += 1;
			}
		}
		System.out.println("Ambiguity Statistics:");
		System.out.println(String.format("A total of %d DS patterns were extracted.", ruleMap.size()));
		System.out.println(String.format("A total of %d were ambiguous.", numAmbig));
		System.out.println(String.format("A total of %d were two-ways ambiguous.", numBiAmbig));
		System.out.println(String.format("A total of %d were three-ways ambiguous.", numTriAmbig));
		System.out.println(String.format("A total of %d were more ambiguous.", numMoreAmbig));
	}
	private static int getTotalExampleCount(UniDSPattern dsPatt, List<UniPSPattern> psPatts, 
			Map<UniRule, List<Example>> ruleExMap) {
		int totalCount = 0;
		for (UniPSPattern psPatt: psPatts) {
			totalCount += ruleExMap.get(new UniRule(dsPatt, psPatt)).size();
		}
		return totalCount;
	}
	public static void printToFile(Map<UniDSPattern, List<UniPSPattern>> ruleMap, 
			Map<UniRule, List<Example>> exMap, String path, String format) {
		printToFile(ruleMap, exMap, path, format, false);
	}
	public static void printToFile(Map<UniDSPattern, List<UniPSPattern>> ruleMap, 
			Map<UniRule, List<Example>> ruleExMap, String path, String format, boolean ambigOnly) {
		try {
			RuleWriter rw = RuleWriter.newWriter(format);
			RuleFileWriter rfw = RuleFileWriter.newWriter(path, format);
			List<UniDSPattern> dsPatts = new LinkedList<UniDSPattern>(ruleMap.keySet());
			Collections.sort(dsPatts);
			if (rfw instanceof RuleAndExampleFileWriter) {
				RuleAndExampleFileWriter refw = (RuleAndExampleFileWriter) rfw;
				String[] includeLayers = new String[] {"PS", "DS-Const"};
				for (UniDSPattern dsPatt: dsPatts) {
					List<UniPSPattern> psPatts = ruleMap.get(dsPatt);
					if (psPatts.size() > 1) {
						int totalRules = psPatts.size();
						int totalExamples = getTotalExampleCount(dsPatt, psPatts, ruleExMap);
						refw.writeHeaderForAmbigRule(rw.writeUnaryDSPattern(dsPatt), totalRules);
						for (UniPSPattern psPatt: psPatts) {
							UniRule rule = new UniRule(dsPatt, psPatt);
							List<Example> exrs = ruleExMap.get(rule);
							int writeLimit = 1;
							if ((float) exrs.size() / totalExamples < 0.33) writeLimit = exrs.size();
							refw.writeRuleAndExampleReps(rw.writeUnaryRule(rule), exrs, includeLayers, writeLimit);
						}
						refw.writeFooterForAmbigRule();
					} else if (psPatts.size() == 1 && ambigOnly == false) {
						UniRule rule = new UniRule(dsPatt, psPatts.get(0));
						List<Example> exrs = ruleExMap.get(rule);
						refw.writeRuleAndExampleReps(rw.writeUnaryRule(rule), exrs, includeLayers, 1);
					}
				}
			}
			else {
				for (UniDSPattern dsPatt: dsPatts) {
					for (UniPSPattern psPatt: ruleMap.get(dsPatt)) {
						UniRule rule = new UniRule(dsPatt, psPatt);
						List<Example> exrs = new LinkedList<Example>(ruleExMap.get(rule));
						rfw.writeRuleRep(rw.writeUnaryRule(rule), exrs.size());
					}
				}
			}
			rfw.close();
		} catch(Exception e) {
			pHandler.handle(e, "Printing of unary rules failed.", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
	}
	
	public static UniRule extractUnaryRule(Token dsToken, Configuration config, boolean onlyMapPOS) {
		
		UniDSPattern dsPatt = ConvertUtils.extractUnaryDSPattern(dsToken, config.getHeadFeats(),
											config.useExtPOS());
		UniPSPattern psPatt;
		if (onlyMapPOS) {
			psPatt = ConvertUtils.extractUnaryPSPattern(dsToken, 1);
		} else {
			psPatt = ConvertUtils.extractUnaryPSPattern(dsToken);
		}
		return new UniRule(dsPatt, psPatt);
	}	
	public static Set<UniRule> extractUnaryRules(Tree tree, Configuration config) {
		return extractUnaryRules(tree, config, false);
	}
	public static Set<UniRule> extractUnaryRules(Tree tree, Configuration config, boolean onlyMapPOS) {
		Set<UniRule> uniRules = new HashSet<UniRule>();
		for(Token t: tree.getTokenList()) {
			// For extracting the POS map we will ignore empty categories
			if (onlyMapPOS && t.isEmptyCategory()) {
				continue;
			}
			uniRules.add(extractUnaryRule(t, config, onlyMapPOS));
		}
		return uniRules;
	}
	private static Map<UniRule, List<Example>> extractRulesFromInputFile(Configuration config) {
		Map<UniRule, List<Example>> ruleExMap = new HashMap<UniRule, List<Example>>();
		int numExamples = 0;
		int numCompat = 0;
		try {
			ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, "xml");
			ExampleDoc exf = exfr.read();
			TreeReader dsr = TreeReader.newReader(opts.dsTreeFormat);
			TreeReader psr = TreeReader.newReader(opts.psTreeFormat);
			Example exr = null;
			while ((exr = exf.getNextExample()) != null) {
				numExamples++;
				// For extracting the POS map we will extract even from incompatible examples
				if (ConvertUtils.isCompatible(exr, opts.onlyMapPOS)) {
					numCompat++;
					try {
						Tree ds = dsr.readTree(exr.getTreeRep("DS-Const"));
						Tree ps = psr.readTree(exr.getTreeRep("PS"));
						Tree tree;
						// For extracting the POS map we will not attempt to merge empty categories
						if (opts.onlyMapPOS) {
							tree = MergeUtils.mergeDSandPSTrees(ds, ps, false, false);
						} else {
							tree = MergeUtils.mergeDSandPSTrees(ds, ps);
						}
						Set<UniRule> uniRules = extractUnaryRules(tree, config, opts.onlyMapPOS);
						for(UniRule uniRule: uniRules) {
							if (!ruleExMap.containsKey(uniRule)) {
								ruleExMap.put(uniRule, new LinkedList<Example>());
							}	
							ruleExMap.get(uniRule).add(exr);
						}
					}
					catch (TreeMergeException e) {
						String errStr = String.format("Could not merge representations for example '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
					catch (UnreadableRepresentationException e) {
						String errStr = String.format("Could not process representation for example '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
				}
			
			}
			System.out.println("Extraction of unary rules was successful.");
			System.out.println(String.format("A total of %d rules were extracted from %d compatible tree pairs.", 
					ruleExMap.size(), numCompat));
			System.out.println(String.format("Processed %d total examples", numExamples));
		} catch(Exception e) {
			pHandler.handle(e, "Extraction of unary rules failed.", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
		return ruleExMap;	
	}

	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			Configuration config;
			if (opts.configFilePath != null) {
				XMLConfiguration xmlConfig = new XMLConfiguration(opts.configFilePath);
				config = Configuration.init(xmlConfig);
			} else {
				config = Configuration.initDefault();
			}
			Map<UniRule, List<Example>> ruleExMap = extractRulesFromInputFile(config);
			Map<UniDSPattern, List<UniPSPattern>> ruleMap = new HashMap<UniDSPattern, List<UniPSPattern>>();
			for (UniRule rule: ruleExMap.keySet()) {
				if (!ruleMap.containsKey(rule.getDSPattern())) {
					ruleMap.put(rule.getDSPattern(), new LinkedList<UniPSPattern>());
				}
				ruleMap.get(rule.getDSPattern()).add(rule.getPSPattern());
			}
			printRuleStatistics(ruleMap);
			printToFile(ruleMap, ruleExMap, opts.outputFilePath, "text");
			if (opts.latexOutputFilePath != null) {
				printToFile(ruleMap, ruleExMap, opts.latexOutputFilePath, "latex", true);
			}
			
		}
		catch(CmdLineException e) {
			System.err.println(e.getMessage());
			System.err.print("Usage: ExtractUnaryRules");
			parser.printSingleLineUsage(System.err);
			System.err.println();
			parser.printUsage(System.err);
			System.exit(1);
		}
		catch(ConfigurationException e) {
			String errStr = String.format("configuration file \'%s\' could not be read.", opts.configFilePath);
			pHandler.handle(e, errStr);
			System.exit(1);
		}
	}
}
